#!/usr/bin/env python3
# -*- coding: utf-8 -*-

from flask import Flask, render_template, request, jsonify, send_file, send_from_directory
from flask_cors import CORS
from flask_socketio import SocketIO, emit
import yaml
import subprocess
import os
import threading
import time
import json
import signal
import psutil
import re
from datetime import datetime
from pathlib import Path
from proxy_server import start_proxy_async
import asyncio

app = Flask(__name__, static_folder='templates/static', static_url_path='/static')
CORS(app)
socketio = SocketIO(app, cors_allowed_origins="*")

# 全局变量
proxy_process_manager = None

class ConfigManager:
    def __init__(self, config_path='config.yaml'):
        self.config_path = config_path
    
    def load_config(self):
        """加载配置"""
        if not os.path.exists(self.config_path):
            return {'rules': []}
        with open(self.config_path, 'r') as f:
            return yaml.safe_load(f)
    
    def save_config(self, config):
        """保存配置"""
        with open(self.config_path, 'w') as f:
            yaml.dump(config, f, default_flow_style=False, allow_unicode=True)
    
    def add_rule(self, rule):
        """添加规则"""
        config = self.load_config()
        config['rules'].append(rule)
        self.save_config(config)
    
    def update_rule(self, index, rule):
        """更新规则"""
        config = self.load_config()
        if 0 <= index < len(config['rules']):
            config['rules'][index] = rule
            self.save_config(config)
    
    def delete_rule(self, index):
        """删除规则"""
        config = self.load_config()
        if 0 <= index < len(config['rules']):
            config['rules'].pop(index)
            self.save_config(config)
    
    def toggle_rule(self, index):
        """切换规则的启用/禁用状态"""
        config = self.load_config()
        if 0 <= index < len(config['rules']):
            rule = config['rules'][index]
            # 切换enabled状态，默认为True
            rule['enabled'] = not rule.get('enabled', True)
            self.save_config(config)
            return rule['enabled']
        return None
    
    def set_rule_enabled(self, index, enabled):
        """设置规则的启用/禁用状态"""
        config = self.load_config()
        if 0 <= index < len(config['rules']):
            config['rules'][index]['enabled'] = enabled
            self.save_config(config)
            return True
        return False

config_manager = ConfigManager()

class SystemProxyManager:
    def __init__(self):
        self.network_service = "Wi-Fi"  # 默认网络服务
    
    def get_network_services(self):
        """获取网络服务列表"""
        try:
            result = subprocess.run(['networksetup', '-listallnetworkservices'], 
                                  capture_output=True, text=True)
            services = []
            for line in result.stdout.split('\n')[1:]:  # 跳过第一行说明
                line = line.strip()
                if line and not line.startswith('*'):
                    services.append(line)
            return services
        except:
            return ["Wi-Fi"]
    
    def get_proxy_status(self):
        """获取代理状态"""
        try:
            http_result = subprocess.run(['networksetup', '-getwebproxy', self.network_service], 
                                       capture_output=True, text=True)
            https_result = subprocess.run(['networksetup', '-getsecurewebproxy', self.network_service], 
                                        capture_output=True, text=True)
            
            http_enabled = "Yes" in http_result.stdout
            https_enabled = "Yes" in https_result.stdout
            
            return {
                'http_enabled': http_enabled,
                'https_enabled': https_enabled,
                'network_service': self.network_service
            }
        except:
            return {'http_enabled': False, 'https_enabled': False, 'network_service': self.network_service}
    
    def enable_proxy(self, host='127.0.0.1', port=8080):
        """启用系统代理"""
        try:
            subprocess.run(['networksetup', '-setwebproxy', self.network_service, host, str(port)])
            subprocess.run(['networksetup', '-setsecurewebproxy', self.network_service, host, str(port)])
            subprocess.run(['networksetup', '-setwebproxystate', self.network_service, 'on'])
            subprocess.run(['networksetup', '-setsecurewebproxystate', self.network_service, 'on'])
            return True
        except:
            return False
    
    def disable_proxy(self):
        """禁用系统代理"""
        try:
            subprocess.run(['networksetup', '-setwebproxystate', self.network_service, 'off'])
            subprocess.run(['networksetup', '-setsecurewebproxystate', self.network_service, 'off'])
            return True
        except:
            return False

system_proxy_manager = SystemProxyManager()

class ProxyProcessManager:
    def __init__(self, config_path='config.yaml'):
        self.config_path = config_path
        self.pid_file = 'proxy_server.pid'
        self.process = None
    
    def start_proxy_process(self, host='0.0.0.0', port=8080, log_level=None):
        """启动代理服务器进程"""
        if self.is_proxy_running():
            return True, "代理服务器已在运行"
        
        try:
            # 启动代理服务器作为子进程
            cmd = [
                'python', 'proxy_server.py',
                '--host', host,
                '--port', str(port),
                '--config', self.config_path
            ]
            
            # 如果指定了日志级别，添加到命令行参数
            if log_level:
                cmd.extend(['--log-level', log_level])
            
            # 创建日志文件
            log_file = open('proxy_server.log', 'w')
            
            self.process = subprocess.Popen(
                cmd,
                stdout=log_file,
                stderr=subprocess.STDOUT,
                start_new_session=True  # 创建新的会话
            )
            
            # 保存PID到文件
            with open(self.pid_file, 'w') as f:
                f.write(str(self.process.pid))
            
            # 等待一小段时间确保进程启动
            time.sleep(0.5)
            
            if self.process.poll() is None:  # 进程仍在运行
                return True, f"代理服务器启动成功，PID: {self.process.pid}"
            else:
                return False, "代理服务器启动失败"
                
        except Exception as e:
            return False, f"启动代理服务器时出错: {str(e)}"
    
    def stop_proxy_process(self):
        """停止代理服务器进程"""
        try:
            pid = self._get_saved_pid()
            if pid:
                try:
                    # 尝试优雅关闭
                    process = psutil.Process(pid)
                    process.terminate()
                    time.sleep(1)
                    
                    # 检查进程是否还在运行
                    if process.is_running():
                        # 强制终止
                        process.kill()
                        time.sleep(0.5)
                    
                    self._cleanup_pid_file()
                    return True, "代理服务器已停止"
                    
                except ProcessLookupError:
                    self._cleanup_pid_file()
                    return True, "代理服务器已停止"
            else:
                return True, "代理服务器未运行"
                
        except Exception as e:
            return False, f"停止代理服务器时出错: {str(e)}"
    
    def is_proxy_running(self):
        """检查代理服务器是否正在运行"""
        pid = self._get_saved_pid()
        if pid:
            try:
                return psutil.pid_exists(pid) and psutil.Process(pid).is_running()
            except:
                self._cleanup_pid_file()
                return False
        return False
    
    def get_proxy_status(self):
        """获取代理服务器状态"""
        if self.is_proxy_running():
            pid = self._get_saved_pid()
            try:
                process = psutil.Process(pid)
                return {
                    'running': True,
                    'pid': pid,
                    'status': process.status(),
                    'cpu_percent': process.cpu_percent(),
                    'memory_info': process.memory_info()._asdict()
                }
            except:
                return {'running': True, 'pid': pid}
        else:
            return {'running': False}
    
    def _get_saved_pid(self):
        """从文件中读取保存的PID"""
        try:
            if os.path.exists(self.pid_file):
                with open(self.pid_file, 'r') as f:
                    return int(f.read().strip())
        except:
            pass
        return None
    
    def _cleanup_pid_file(self):
        """清理PID文件"""
        try:
            if os.path.exists(self.pid_file):
                os.remove(self.pid_file)
        except:
            pass
    
    def restart_proxy_process(self, host='0.0.0.0', port=8080, log_level=None):
        """重启代理服务器进程"""
        if self.is_proxy_running():
            # 先停止现有进程
            stop_success, stop_message = self.stop_proxy_process()
            if not stop_success:
                return False, f"停止代理服务器失败: {stop_message}"
            
            # 等待进程完全停止
            time.sleep(1)
        
        # 启动新进程
        return self.start_proxy_process(host, port, log_level)

# 初始化进程管理器
proxy_process_manager = ProxyProcessManager()

class LogManager:
    def __init__(self, log_file='proxy_server.log'):
        self.log_file = log_file
    
    def parse_log_line(self, line):
        """解析单行日志"""
        line = line.strip()
        if not line:
            return None
        
        # 匹配不同格式的日志
        patterns = [
            # 格式1: 2025-06-22 15:45:55 INFO: 启动代理服务器 0.0.0.0:8080
            (r'^(\d{4}-\d{2}-\d{2} \d{2}:\d{2}:\d{2}) (\w+): (.+)$', 'standard'),
            # 格式2: [15:45:55.359] 加载了 3 条代理规则
            (r'^\[(\d{2}:\d{2}:\d{2}\.\d+)\] (.+)$', 'mitmproxy'),
            # 格式3: 127.0.0.1:56481: GET http://127.0.0.1:31001/xxx
            (r'^(\d+\.\d+\.\d+\.\d+:\d+): (.+)$', 'request'),
            # 格式4:  << [Errno 61] Connect call failed
            (r'^ << (.+)$', 'error'),
        ]
        
        for pattern, log_type in patterns:
            match = re.match(pattern, line)
            if match:
                if log_type == 'standard':
                    return {
                        'timestamp': match.group(1),
                        'level': match.group(2),
                        'message': match.group(3),
                        'type': log_type
                    }
                elif log_type == 'mitmproxy':
                    # 转换时间格式
                    time_str = match.group(1)
                    today = datetime.now().strftime('%Y-%m-%d')
                    timestamp = f"{today} {time_str.split('.')[0]}"
                    return {
                        'timestamp': timestamp,
                        'level': 'INFO',
                        'message': match.group(2),
                        'type': log_type
                    }
                elif log_type == 'request':
                    today = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
                    return {
                        'timestamp': today,
                        'level': 'DEBUG',
                        'message': f"[{match.group(1)}] {match.group(2)}",
                        'type': log_type
                    }
                elif log_type == 'error':
                    today = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
                    return {
                        'timestamp': today,
                        'level': 'ERROR',
                        'message': match.group(1),
                        'type': log_type
                    }
        
        # 如果都不匹配，作为普通信息处理
        today = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
        return {
            'timestamp': today,
            'level': 'INFO',
            'message': line,
            'type': 'unknown'
        }
    
    def read_logs(self, rule_filter=None, level_filter=None, date_filter=None, limit=500):
        """读取并过滤日志"""
        logs = []
        
        if not os.path.exists(self.log_file):
            return logs
        
        try:
            with open(self.log_file, 'r', encoding='utf-8') as f:
                lines = f.readlines()
            
            # 解析日志行
            for line in lines:
                log_entry = self.parse_log_line(line)
                if log_entry:
                    logs.append(log_entry)
            
            # 应用过滤器
            filtered_logs = []
            for log in logs:
                # 规则过滤
                if rule_filter and rule_filter not in log['message']:
                    continue
                
                # 级别过滤
                if level_filter and log['level'] != level_filter:
                    continue
                
                # 日期过滤
                if date_filter:
                    try:
                        log_date = log['timestamp'].split(' ')[0]
                        if log_date != date_filter:
                            continue
                    except:
                        continue
                
                filtered_logs.append(log)
            
            # 限制数量并倒序（最新的在前）
            return filtered_logs[-limit:][::-1]
            
        except Exception as e:
            print(f"读取日志失败: {e}")
            return []

log_manager = LogManager()

@app.route('/')
def index():
    """主页"""
    return render_template('index.html')



@app.route('/api/config', methods=['GET'])
def get_config():
    """获取配置"""
    config = config_manager.load_config()
    return jsonify(config)

@app.route('/api/config', methods=['POST'])
def save_config():
    """保存配置"""
    config = request.json
    config_manager.save_config(config)
    return jsonify({'success': True})

@app.route('/api/rules', methods=['POST'])
def add_rule():
    """添加规则"""
    rule = request.json
    config_manager.add_rule(rule)
    
    # 如果代理服务器正在运行，重启以应用新配置
    if proxy_process_manager.is_proxy_running():
        restart_success, restart_message = proxy_process_manager.restart_proxy_process()
        return jsonify({
            'success': True, 
            'restarted': restart_success,
            'restart_message': restart_message
        })
    
    return jsonify({'success': True, 'restarted': False})

@app.route('/api/rules/<int:index>', methods=['PUT'])
def update_rule(index):
    """更新规则"""
    rule = request.json
    config_manager.update_rule(index, rule)
    
    # 如果代理服务器正在运行，重启以应用新配置
    if proxy_process_manager.is_proxy_running():
        restart_success, restart_message = proxy_process_manager.restart_proxy_process()
        return jsonify({
            'success': True, 
            'restarted': restart_success,
            'restart_message': restart_message
        })
    
    return jsonify({'success': True, 'restarted': False})

@app.route('/api/rules/<int:index>', methods=['DELETE'])
def delete_rule(index):
    """删除规则"""
    config_manager.delete_rule(index)
    
    # 如果代理服务器正在运行，重启以应用新配置
    if proxy_process_manager.is_proxy_running():
        restart_success, restart_message = proxy_process_manager.restart_proxy_process()
        return jsonify({
            'success': True, 
            'restarted': restart_success,
            'restart_message': restart_message
        })
    
    return jsonify({'success': True, 'restarted': False})

@app.route('/api/rules/<int:index>/toggle', methods=['POST'])
def toggle_rule(index):
    """切换规则的启用/禁用状态"""
    new_status = config_manager.toggle_rule(index)
    if new_status is None:
        return jsonify({'success': False, 'error': '规则不存在'}), 404
    
    # 如果代理服务器正在运行，重启以应用新配置
    if proxy_process_manager.is_proxy_running():
        restart_success, restart_message = proxy_process_manager.restart_proxy_process()
        return jsonify({
            'success': True,
            'enabled': new_status,
            'restarted': restart_success,
            'restart_message': restart_message
        })
    
    return jsonify({'success': True, 'enabled': new_status, 'restarted': False})

@app.route('/api/rules/<int:index>/enable', methods=['POST'])
def enable_rule(index):
    """启用规则"""
    data = request.json or {}
    enabled = data.get('enabled', True)
    
    success = config_manager.set_rule_enabled(index, enabled)
    if not success:
        return jsonify({'success': False, 'error': '规则不存在'}), 404
    
    # 如果代理服务器正在运行，重启以应用新配置
    if proxy_process_manager.is_proxy_running():
        restart_success, restart_message = proxy_process_manager.restart_proxy_process()
        return jsonify({
            'success': True,
            'enabled': enabled,
            'restarted': restart_success,
            'restart_message': restart_message
        })
    
    return jsonify({'success': True, 'enabled': enabled, 'restarted': False})

@app.route('/api/proxy/status', methods=['GET'])
def get_proxy_status():
    """获取代理状态"""
    proxy_status = proxy_process_manager.get_proxy_status()
    return jsonify({
        'server_running': proxy_status['running'],
        'server_details': proxy_status,
        'system_proxy': system_proxy_manager.get_proxy_status()
    })

@app.route('/api/proxy/start', methods=['POST'])
def start_proxy_server():
    """启动代理服务器"""
    data = request.json or {}
    host = data.get('host', '0.0.0.0')
    port = data.get('port', 8080)
    log_level = data.get('log_level')  # 可选的日志级别参数
    
    success, message = proxy_process_manager.start_proxy_process(host, port, log_level)
    return jsonify({
        'success': success,
        'message': message,
        'running': proxy_process_manager.is_proxy_running()
    })

@app.route('/api/proxy/stop', methods=['POST'])
def stop_proxy_server():
    """停止代理服务器"""
    success, message = proxy_process_manager.stop_proxy_process()
    return jsonify({
        'success': success,
        'message': message,
        'running': proxy_process_manager.is_proxy_running()
    })

@app.route('/api/system-proxy/enable', methods=['POST'])
def enable_system_proxy():
    """启用系统代理"""
    data = request.json
    host = data.get('host', '127.0.0.1')
    port = data.get('port', 8080)
    success = system_proxy_manager.enable_proxy(host, port)
    return jsonify({'success': success})

@app.route('/api/system-proxy/disable', methods=['POST'])
def disable_system_proxy():
    """禁用系统代理"""
    success = system_proxy_manager.disable_proxy()
    return jsonify({'success': success})

@app.route('/api/network-services', methods=['GET'])
def get_network_services():
    """获取网络服务列表"""
    services = system_proxy_manager.get_network_services()
    return jsonify({'services': services})

@app.route('/api/logs', methods=['GET'])
def get_logs():
    """获取日志列表"""
    rule_filter = request.args.get('rule', '')
    level_filter = request.args.get('level', '')
    date_filter = request.args.get('date', '')
    limit = int(request.args.get('limit', 500))
    
    logs = log_manager.read_logs(
        rule_filter=rule_filter if rule_filter else None,
        level_filter=level_filter if level_filter else None,
        date_filter=date_filter if date_filter else None,
        limit=limit
    )
    
    return jsonify({
        'logs': logs,
        'total': len(logs),
        'filters': {
            'rule': rule_filter,
            'level': level_filter,
            'date': date_filter
        }
    })

@app.route('/api/cert/status', methods=['GET'])
def get_cert_status():
    """获取证书状态"""
    try:
        cert_path = Path.home() / '.mitmproxy' / 'mitmproxy-ca-cert.pem'
        
        if cert_path.exists():
            # 获取证书创建时间
            stat = cert_path.stat()
            created_at = datetime.fromtimestamp(stat.st_ctime).strftime('%Y-%m-%d %H:%M:%S')
            
            return jsonify({
                'success': True,
                'exists': True,
                'path': str(cert_path),
                'created_at': created_at
            })
        else:
            return jsonify({
                'success': True,
                'exists': False,
                'path': str(cert_path)
            })
    except Exception as e:
        logger.error(f"检查证书状态失败: {e}")
        return jsonify({'success': False, 'error': str(e)}), 500

@app.route('/api/cert/download', methods=['GET'])
def download_cert():
    """下载证书文件"""
    try:
        cert_path = Path.home() / '.mitmproxy' / 'mitmproxy-ca-cert.pem'
        
        if not cert_path.exists():
            return jsonify({'success': False, 'error': '证书文件不存在'}), 404
        
        return send_file(
            cert_path,
            as_attachment=True,
            download_name='mitmproxy-ca-cert.pem',
            mimetype='application/x-pem-file'
        )
    except Exception as e:
        print(f"下载证书失败: {e}")
        return jsonify({'success': False, 'error': str(e)}), 500

if __name__ == '__main__':
    socketio.run(app, host='0.0.0.0', port=5001, debug=True) 